from tensorflow.keras.datasets import mnist
from tensorflow import keras
from tensorflow.keras import layers
from matplotlib import pyplot


def main():
    (tr_images, tr_labels), (te_images, te_labels) = mnist.load_data()
    print('train shape: %s', tr_images.shape)
    print('train labels: %s', tr_labels)

    net = keras.Sequential([layers.Dense(512, activation='relu'), layers.Dense(10, activation='softmax')])
    net.compile(optimizer='rmsprop', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

    pyplot.imshow(tr_images[0], cmap=pyplot.cm.binary)
    pyplot.show()

    tr_images = tr_images.reshape((60000, 28 * 28)).astype('float32') / 255
    te_images = te_images.reshape((10000, 28 * 28)).astype('float32') / 255

    net.fit(tr_images, tr_labels, epochs=5, batch_size=128)

    te_pred = net.predict(te_images[:10])
    print([i.argmax() for i in te_pred])

    te_loss, te_acc = net.evaluate(te_images, te_labels)
    print('test accuracy: %s' % te_acc)



if __name__ == '__main__':
    main()
